import json
from itertools import chain

from tqdm import tqdm
import pandas as pd
import torch
from torch_geometric.loader import DataLoader
from dataloader import get_load_dataset
from extract_data import extract_test_graph_data
from model import PredictModel

import const as cn
from utils import load_pickle, eval_batch


def load_best_epoch(_type):
    print('loading best epoch')
    return torch.load(cn.FILE_BEST_MODEL_MODEL.format(type=_type), map_location='cpu')


def eval_test(model, device, loader):
    model.eval()
    y_pred = []
    for _, batch in enumerate(tqdm(loader, desc="Iteration")):
        eval_batch(batch, device, model, y_pred)
    return list(chain(*y_pred))


def run(device, net_parameters, _type, x_test):
    type_word2id = load_pickle(cn.FILE_WORD_TYPES.format(type=_type))
    value_word2id = load_pickle(cn.FILE_WORD_VALUES.format(type=_type))
    net_parameters["type_nums"] = len(type_word2id)
    net_parameters["value_nums"] = len(value_word2id)

    extract_test_graph_data(_type, x_test, cn.DIR_TEST_AST, cn.DIR_TMP_TEST_DATASET, type_word2id, value_word2id)

    load_test_dataset = get_load_dataset(cn.DIR_TMP_TEST_DATASET, "test", x_test)
    test_loader = DataLoader(load_test_dataset(), batch_size=net_parameters["batch_size"])
    model_state_dict, _ = load_best_epoch(_type)

    model = PredictModel(net_parameters).to(device)

    model.load_state_dict(model_state_dict)
    model = model.to(device)
    test_result = eval_test(model, device, test_loader)

    return test_result


if __name__ == "__main__":
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    with open(cn.FILE_CONFIG_JSON) as f:
        net_parameters = json.load(f)

    x_test = pd.read_csv(cn.FILE_TEST_CSV_DATA)
    out = open(cn.FILE_OUTPUT_CSV, "w+")
    out.write("file_id,type\n")

    for _type in ('jsp', 'php'):
        x_t_test = x_test[x_test['type'] == _type]
        print("type: {}  length: {}", _type, x_t_test.shape)
        ex_flag = 0
        try:
            test_result = run(device, net_parameters, _type, x_t_test)
        except Exception as err:
            print(err)
            ex_flag = 1
        data_list = x_t_test['file_id'].tolist()

        if not ex_flag:
            for index in range(len(test_result)):
                file_id = data_list[index]
                label = "white" if test_result[index].item() == 0 else "black"
                out.write("{},{}\n".format(file_id, label))
            continue

        for file_id in data_list:
            out.write("{},black\n".format(file_id))

    out.close()
